
from langchain.retrievers.self_query.base import SelfQueryRetriever
from langchain.chains.query_constructor.base import AttributeInfo
from langchain_openai import AzureChatOpenAI
from langchain_community.vectorstores import Chroma
from langchain_openai.embeddings import AzureOpenAIEmbeddings  # 导入嵌入模型
from tool import get_azure_endpoint,get_api_version,get_api_key

if __name__ == '__main__':
    llm = AzureChatOpenAI(
        azure_endpoint=get_azure_endpoint().rstrip('/'),  # 移除尾部斜杠，只保留基础URL
        azure_deployment="gpt-4o-mini",  # 重命名为 azure_deployment
        openai_api_version=get_api_version(),  # 参数名不变
        openai_api_key=get_api_key(),
        openai_api_type="azure",
    )

    embedding = AzureOpenAIEmbeddings(
        azure_endpoint=get_azure_endpoint().rstrip('/'),  # 移除尾部斜杠，只保留基础URL
        model="text-embedding-3-small",  # 重命名为 azure_deployment
        api_key=get_api_key(),
        api_version=get_api_version()
    )
    vectordb_chinese = Chroma(
        persist_directory="./docs/chroma/matplotlib",
        embedding_function=embedding,
    )

    metadata_field_info_chinese = [
        AttributeInfo(
            name="source",
            description="The lecture the chunk is from, should be one of `./docs/matplotlib/第一回：Matplotlib初相识.pdf`, `./docs/matplotlib/第二回：艺术画笔见乾坤.pdf`, or `./docs/matplotlib/第三回：布局格式定方圆.pdf`, or `./docs/matplotlib/第四回：进阶绘图技巧.pdf`, or `./docs/matplotlib/第五回：样式色彩秀芳华.pdf`",
            type="string",
        ),
        AttributeInfo(
            name="page",
            description="The page from the lecture",
            type="integer",
        ),
    ]
    document_content_description_chinese = "Matplotlib 课堂讲义"
    retriever_chinese = SelfQueryRetriever.from_llm(
        llm,
        vectordb_chinese,
        document_content_description_chinese,
        metadata_field_info_chinese,
        verbose=True
    )
    question_chinese = "他们在第二讲中对Figure做了些什么?"

    docs_chinese = retriever_chinese.get_relevant_documents(question_chinese)
    for d in docs_chinese:
        print(d.metadata)